package org.nimi317.web_gis.model;

import jakarta.validation.constraints.Min;
import org.deeplearning4j.nn.multilayer.MultiLayerNetwork;
import org.deeplearning4j.util.ModelSerializer;
import org.nd4j.linalg.api.ndarray.INDArray;
import org.nd4j.linalg.dataset.api.preprocessor.DataNormalization;
import org.nd4j.linalg.dataset.api.preprocessor.NormalizerMinMaxScaler;
import org.nd4j.linalg.dataset.api.preprocessor.NormalizerStandardize;
import org.nd4j.linalg.factory.Nd4j;
import org.nd4j.linalg.indexing.INDArrayIndex;
import org.nd4j.linalg.indexing.NDArrayIndex;
import org.nimi317.web_gis.entity.Model;
import org.nimi317.web_gis.exception.E;
import org.nimi317.web_gis.exception.RException;
import org.nimi317.web_gis.runnable.ModelRunnable;

import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;

/**
 * @author thunderobot
 */
public class PredictModel {

    private final MultiLayerNetwork model;

    private final DataNormalization normalizer;

    private final INDArray featureMask;

    private final INDArray labelMask;

    private final int featureRange;

    private final int labelRange;

    private final int featureStep;

    private final int labelStep;

    public PredictModel(String path, boolean isStandardize, int featureColumnsStart, int featureColumnsEnd, int labelColumnsStart, int labelColumnsEnd, int featureStep, int labelStep) {
        // 检测path是否存在
        Path base = Paths.get(ModelRunnable.baseLocation);
        Path directory = base.resolve(path);
        if (!Files.exists(directory)) {
            throw new RException(E.ModelNotFoundException);
        }
        this.featureRange = featureColumnsEnd - featureColumnsStart;
        this.labelRange = featureColumnsEnd - featureColumnsStart;
        this.featureStep = featureStep;
        this.labelStep = labelStep;
        int maxStep = Math.max(featureStep, labelStep);
        try {
            // 掩膜数组
            INDArray zeros = Nd4j.zeros(1, maxStep);
            INDArray zeros1 = Nd4j.zeros(1, maxStep);
            zeros.put(new INDArrayIndex[]{NDArrayIndex.all(), NDArrayIndex.interval(0, labelStep)}, Nd4j.ones(1, labelStep));
            zeros1.put(new INDArrayIndex[]{NDArrayIndex.all(), NDArrayIndex.interval(0, featureStep)}, Nd4j.ones(1, featureStep));

            this.featureMask = zeros1;
            this.labelMask = zeros;
            this.model = ModelSerializer.restoreMultiLayerNetwork(directory.resolve("model.zip").normalize().toFile());
            if (isStandardize) {
                NormalizerStandardize standardize = new NormalizerStandardize();
                standardize.fitLabel(true);
                standardize.load(
                        directory.resolve("mean.bin").toFile(),
                        directory.resolve("std.bin").toFile(),
                        directory.resolve("label_mean.bin").toFile(),
                        directory.resolve("label_std.bin").toFile()
                );
                this.normalizer = standardize;
            } else {
                NormalizerMinMaxScaler scaler = new NormalizerMinMaxScaler();
                scaler.fitLabel(true);
                scaler.load(
                        directory.resolve("min.bin").toFile(),
                        directory.resolve("max.bin").toFile(),
                        directory.resolve("label_min.bin").toFile(),
                        directory.resolve("label_max.bin").toFile()
                );
                this.normalizer = scaler;
            }
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    public PredictModel(Model model) {
        this(model.getUrl(), model.getStandard(), model.getFeatureStart(), model.getFeatureStop(), model.getLabelStart(), model.getLabelStop(), model.getFeatureStep(), model.getLabelStep());
    }

    /**
     * 单步预测
     *
     * @param input 输入
     * @return 输出
     */
    public double[][] predict(double[][] input) {
        check(input);
        try (INDArray transpose = Nd4j.create(input).transpose()) {
            //[用例,列数,行数]
            int rows = transpose.rows();
            int columns = transpose.columns();
            INDArray reshape = transpose.reshape(1, rows, columns);
            normalizer.transform(reshape);
            INDArray output = this.model.output(reshape, false, featureMask, labelMask);
            normalizer.revertLabels(output, labelMask);
            return output.reshape(rows, columns).transpose().toDoubleMatrix();
        }
    }

    public double[][] predictMul(double[][] input) {
//        check(input);
        try (INDArray transpose = Nd4j.create(input).transpose()) {
            //[用例,列数,行数]
            int rows = transpose.rows();
            int columns = transpose.columns();
            INDArray reshape = transpose.reshape(1, rows, columns);
            normalizer.transform(reshape);
            INDArray output = this.model.output(reshape, false, featureMask, labelMask);
            //[num,step]
            normalizer.revertLabels(output, labelMask);
            return output.reshape(rows, columns).toDoubleMatrix();
        }
    }

    /**
     * 多步预测
     *
     * @param input 输入
     * @return 输出
     */
    public double[][][] predict(double[][] input, @Min(value = 1) int steps) {
        check(input);
        try (INDArray transpose = Nd4j.create(input).transpose()) {
            //[用例,列数,行数]
            int rows = transpose.rows();
            int columns = transpose.columns();
            INDArray reshape = transpose.reshape(1, rows, columns);
            normalizer.transform(reshape);
            double[][][] result = new double[steps][rows][columns];
            for (int i = 0; i < steps; i++) {
                INDArray step = this.model.rnnTimeStep(reshape);
                normalizer.revertLabels(step, labelMask);
                double[][] floats = step.reshape(step.size(1), step.size(2)).transpose().toDoubleMatrix();
                result[i] = floats;
            }
            return result;
        }
    }

    public void close() {
        this.model.close();
    }

    private void check(double[][] input) {
        //获取纬度度信息
        int rows = input.length;
        int columns = input[0].length;

        if (rows != this.featureStep || columns != this.featureRange) {
            throw new RException(E.PredictInputNotMatch);
        }
    }
}
